#include "dimp/dimp_optimizer.hpp"
#include <iostream>
#include "dimp/model_process.h"
#include "acl/acl.h"
#include "utils.h"

using namespace std;
using namespace cv;

Optimizer::Optimizer(const char* model_01_path, const char* model_02_path)
:deviceId_(0), context_(nullptr), stream_(nullptr), isInited_(false){
    initial_optimizer_path_ = model_01_path;
    tracking_optimizer_path_ = model_02_path;
}

Optimizer::~Optimizer(){
    DestroyResource();
}

Result Optimizer::Init(int32_t& deviceId, aclrtContext& context, aclrtStream& stream, aclrtRunMode& runMode) {
    //If it is already initialized, it is returned
    if (isInited_) {
        INFO_LOG("Classify instance is initied already!");
        return SUCCESS;
    }


    deviceId_ = deviceId;
    context_ = context;
    stream_ = stream;
    runMode_ = runMode;

    //Initializes the model management instance
    AtlasError ret = InitModel();
    if (ret != SUCCESS) {
        ERROR_LOG("Init model failed");
        return FAILED;
    }
    
    isInited_ = true;

    ret = CreateInputdDataset_for_InitialOptimizer();
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("CreateInputdDataset_for_InitialOptimizer failed!");
    }

    ret = CreateInputdDataset_for_TrackingOptimizer();
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("CreateInputdDataset_for_TrackingOptimizer failed!");
    }


    return SUCCESS;

}

Result Optimizer::InitModel() {
    Result ret = initial_optimizer_.LoadModelFromFileWithMem(initial_optimizer_path_);
    if (ret != SUCCESS) {
        ERROR_LOG("execute LoadModelFromFileWithMem failed");
        return FAILED;
    }

    ret = initial_optimizer_.CreateDesc();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateDesc failed");
        return FAILED;
    }

    ret = initial_optimizer_.CreateOutput();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateOutput failed");
        return FAILED;
    }

    ret = tracking_optimizer_.LoadModelFromFileWithMem(tracking_optimizer_path_);
    if (ret != SUCCESS) {
        ERROR_LOG("execute LoadModelFromFileWithMem failed");
        return FAILED;
    }

    ret = tracking_optimizer_.CreateDesc();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateDesc failed");
        return FAILED;
    }

    ret = tracking_optimizer_.CreateOutput();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateOutput failed");
        return FAILED;
    }


    return SUCCESS;
}

Result Optimizer::CreateInputdDataset_for_InitialOptimizer()
{
    aclError aclRet;
    aclRet = aclrtMalloc(&init_weight_DataBuf_, 16384, ACL_MEM_MALLOC_HUGE_FIRST);
    if (aclRet != ACL_ERROR_NONE) {

        ERROR_LOG("malloc weightsData buffer failed, aclRet is %d", aclRet);
        return FAILED;
    }

     //Request image data memory for input model
    aclRet = aclrtMalloc(&classifier_feats_DataBuf_, 4313088, ACL_MEM_MALLOC_HUGE_FIRST);
    //  cout << imageDataSize_<< endl;
    if (aclRet != ACL_ERROR_NONE) {

        ERROR_LOG("malloc classifier_feats buffer failed, aclRet is %d", aclRet);
        return FAILED;
    }

    bboxesSize_ = sizeof(bboxes_);
    if (runMode_ == ACL_HOST)
        bboxesBuf_ = Utils::CopyDataHostToDevice((void *)bboxes_, bboxesSize_);
    else
        bboxesBuf_ = Utils::CopyDataDeviceToDevice((void *)bboxes_, bboxesSize_);
    if (bboxesBuf_ == nullptr) {
        ERROR_LOG("Copy roiboxes to device failed");
        return FAILED;
    }

    

    //Use the applied memory to create the model and input dataset. After creation, only update the memory data for each frame of inference, instead of creating the input dataset every time
    Result ret = initial_optimizer_.Create3Inputs(init_weight_DataBuf_, 16384, classifier_feats_DataBuf_, 4313088, bboxesBuf_, bboxesSize_);
    if (ret != SUCCESS) {
        ERROR_LOG("Create mode input dataset failed");
        return FAILED;
    }

    return SUCCESS;
}

Result Optimizer::CreateInputdDataset_for_TrackingOptimizer()
{
    aclError aclRet;
    aclRet = aclrtMalloc(&origin_weight_DataBuf_, 16384, ACL_MEM_MALLOC_HUGE_FIRST);
    if (aclRet != ACL_ERROR_NONE) {

        ERROR_LOG("malloc weightsData buffer failed, aclRet is %d", aclRet);
        return FAILED;
    }

     //Request image data memory for input model
    aclRet = aclrtMalloc(&tracking_classifier_feats_DataBuf_, 16588800, ACL_MEM_MALLOC_HUGE_FIRST);
    //  cout << imageDataSize_<< endl;
    if (aclRet != ACL_ERROR_NONE) {

        ERROR_LOG("malloc classifier_feats buffer failed, aclRet is %d", aclRet);
        return FAILED;
    }

    tracking_bboxesSize_ = sizeof(tracking_bboxes_);
    if (runMode_ == ACL_HOST)
        tracking_bboxesBuf_ = Utils::CopyDataHostToDevice((void *)tracking_bboxes_, tracking_bboxesSize_);
    else
        tracking_bboxesBuf_ = Utils::CopyDataDeviceToDevice((void *)tracking_bboxes_, tracking_bboxesSize_);
    if (tracking_bboxesBuf_ == nullptr) {
        ERROR_LOG("Copy tracking_bboxes to device failed");
        return FAILED;
    }

    sample_weightSize_ = sizeof(sample_weight_);
    if (runMode_ == ACL_HOST)
        sample_weightBuf_ = Utils::CopyDataHostToDevice((void *)sample_weight_, sample_weightSize_);
    else
        sample_weightBuf_ = Utils::CopyDataDeviceToDevice((void *)sample_weight_, sample_weightSize_);
    if (sample_weightBuf_ == nullptr) {
        ERROR_LOG("Copy sample_weight to device failed");
        return FAILED;
    }

    //Use the applied memory to create the model and input dataset. After creation, only update the memory data for each frame of inference, instead of creating the input dataset every time
    Result ret = tracking_optimizer_.Create4Inputs(origin_weight_DataBuf_, 16384, tracking_classifier_feats_DataBuf_, 16588800, tracking_bboxesBuf_, tracking_bboxesSize_, sample_weightBuf_, sample_weightSize_);
    if (ret != SUCCESS) {
        ERROR_LOG("Create mode input dataset failed");
        return FAILED;
    }

    return SUCCESS;
}


Result Optimizer::LoadInputDataForInitialOptimizer(Nodes classifier_feats, Nodes init_weight, vector<Target_Box> transform_boxes) {
    aclrtMemcpyKind policy = (runMode_ == ACL_HOST)?
                             ACL_MEMCPY_HOST_TO_DEVICE:ACL_MEMCPY_DEVICE_TO_DEVICE;

    Result ret;

    // float *weight = (float *)(init_weight.firstaddr);
    // for(int i = 0; i < 200; ++i)
    // {
    //     cout<<"init_weight = "<<weight[i]<<endl;
    // }

    ret = (Result)aclrtMemcpy(init_weight_DataBuf_, 16384,
                            init_weight.firstaddr, 16384, policy);
    
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("Copy initweight to device failed.");
        return FAILED;
    }

    // float *classifier = (float *)(classifier_feats.firstaddr);
    // for(int i = 0; i < 200; ++i)
    // {
    //     cout<<"classifier = "<<classifier[i]<<endl;
    // }

    ret = (Result)aclrtMemcpy(classifier_feats_DataBuf_, 4313088,
                            classifier_feats.firstaddr, 4313088, policy);
    
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("Copy classifier_feats to device failed.");
        return FAILED;
    }

    for (int i = 0; i < 13; ++i){
        bboxes_[i][0] = transform_boxes[i].pos.x;
        bboxes_[i][1] = transform_boxes[i].pos.y;
        bboxes_[i][2] = transform_boxes[i].size.x;
        bboxes_[i][3] = transform_boxes[i].size.y;
    }

    // bboxesSize_ = sizeof(bboxes_);
    // if (runMode_ == ACL_HOST)
    //     bboxesBuf_ = Utils::CopyDataHostToDevice((void *)bboxes_, bboxesSize_);
    // else
    //     bboxesBuf_ = Utils::CopyDataDeviceToDevice((void *)bboxes_, bboxesSize_);
    // if (bboxesBuf_ == nullptr) {
    //     ERROR_LOG("Copy roiboxes to device failed");
    //     return FAILED;
    // }

    ret = (Result)aclrtMemcpy(bboxesBuf_, bboxesSize_,
                            bboxes_, bboxesSize_, policy);
    
    if (ret != ACL_ERROR_NONE) {
          ERROR_LOG("Copy bbox to device failed.");
          return FAILED;
    }


    return SUCCESS;

}

Result Optimizer::LoadInputDataForTrackingOptimizer(Nodes classifier_feats, void* origin_weight, int train_nums ,int frame_nums, int now_state, Vector now_pos, Vector now_size) {
    aclrtMemcpyKind policy = (runMode_ == ACL_HOST)?
                             ACL_MEMCPY_HOST_TO_DEVICE:ACL_MEMCPY_DEVICE_TO_DEVICE;

    Result ret;

    // float *weight = (float *)(init_weight.firstaddr);
    // for(int i = 0; i < 200; ++i)
    // {
    //     cout<<"init_weight = "<<weight[i]<<endl;
    // }
    // cout<<"begin!"<<endl;
    if (train_nums == 1) {
        for (int i = 0; i < 50; ++i) {
            if (i < 13) {
                sample_weight_[i] = (1.0 / 13.0) / 0.99 / 1.02; 
            }
            if (i == 13) {
                sample_weight_[i] = 0.01 / 1.02;
            }
           
            if (i > 13) {
                sample_weight_[i] = 0;
            }
        }
        prev_ind_ = 13;

        ret = (Result)aclrtMemcpy(tracking_classifier_feats_DataBuf_, 16588800,
                            classifier_feats_DataBuf_, 4313088, policy);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy classifier_feats to device failed.");
            return FAILED;
        }
        ret = (Result)aclrtMemcpy(tracking_classifier_feats_DataBuf_ + (train_nums + 12) * 331776, 16588800,
                            classifier_feats.firstaddr, 331776, policy);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy classifier_feats to device failed.");
            return FAILED;
        }

        tracking_bboxes_[train_nums + 12][0] = now_pos.x;
        tracking_bboxes_[train_nums + 12][1] = now_pos.y;
        tracking_bboxes_[train_nums + 12][2] = now_size.x;
        tracking_bboxes_[train_nums + 12][3] = now_size.y;

        ret = (Result)aclrtMemcpy(tracking_bboxesBuf_, tracking_bboxesSize_,
                            tracking_bboxes_, tracking_bboxesSize_, policy);
    
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy bbox to device failed.");
            return FAILED;
        }

        ret = (Result)aclrtMemcpy(sample_weightBuf_, sample_weightSize_,
                            sample_weight_, sample_weightSize_, policy);
    
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy sample_weight_ to device failed.");
            return FAILED;
        }
    }
    
    else if (((train_nums > 1) && (train_nums < 18)))
    {
        ret = (Result)aclrtMemcpy(tracking_classifier_feats_DataBuf_ + (train_nums + 12) * 331776, 16588800,
                            classifier_feats.firstaddr, 331776, policy);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy classifier_feats to device failed.");
        return FAILED;
        }

        tracking_bboxes_[train_nums + 12][0] = now_pos.x;
        tracking_bboxes_[train_nums + 12][1] = now_pos.y;
        tracking_bboxes_[train_nums + 12][2] = now_size.x;
        tracking_bboxes_[train_nums + 12][3] = now_size.y;

        ret = (Result)aclrtMemcpy(tracking_bboxesBuf_, tracking_bboxesSize_,
                            tracking_bboxes_, tracking_bboxesSize_, policy);
    
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy bbox to device failed.");
            return FAILED;
        }
        
        float sum = 0.0;
        for (int i = 0; i < 50; ++i) {
            if (i < 13) {
                sum += sample_weight_[i];
            }
            
            else if (i == prev_ind_ + 1) {
                sample_weight_[i] = sample_weight_[prev_ind_] / 0.99;
                sum += sample_weight_[i];
            }
            else {
                sum += sample_weight_[i];
            }
        }
        for (int i = 0; i < 50; ++i) {
            sample_weight_[i] /= sum;
        }
        prev_ind_ = train_nums + 12;

        ret = (Result)aclrtMemcpy(sample_weightBuf_, sample_weightSize_,
                            sample_weight_, sample_weightSize_, policy);
    
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy sample_weight_ to device failed.");
            return FAILED;
        }
    }

    else {
        int replace_id = train_nums % 37;
        
        if (replace_id == 0) {
            replace_id = 37;
        }

        // cout << "replace_id = "<<replace_id<<endl;
        
        ret = (Result)aclrtMemcpy(tracking_classifier_feats_DataBuf_ + (replace_id + 12) * 331776, 16588800,
                            classifier_feats.firstaddr, 331776, policy);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy classifier_feats to device failed.");
        return FAILED;
        }
        // cout<<"copy success!"<<endl;

        tracking_bboxes_[replace_id + 12][0] = now_pos.x;
        tracking_bboxes_[replace_id + 12][1] = now_pos.y;
        tracking_bboxes_[replace_id + 12][2] = now_size.x;
        tracking_bboxes_[replace_id + 12][3] = now_size.y;

        ret = (Result)aclrtMemcpy(tracking_bboxesBuf_, tracking_bboxesSize_,
                            tracking_bboxes_, tracking_bboxesSize_, policy);
    
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy bbox to device failed.");
            return FAILED;
        }
        // cout<<"copy success!"<<endl;

        float sum = 0.0;
        float sum_init = 0.0;
        float sum_else = 0.0;
        // cout<<"prev_ind_ = "<<prev_ind_<<endl;
        for (int i = 0; i < 50; ++i) {
            if (i < 13) {
                sum += sample_weight_[i];
            }

            else if (prev_ind_ == 49 && i == prev_ind_) {
                sum -= sample_weight_[13];
                sample_weight_[13] = sample_weight_[prev_ind_] / 0.99;
                sum = sum + sample_weight_[i] + sample_weight_[13];
            }
            else if (i == prev_ind_ && prev_ind_ != 49) {
                sample_weight_[prev_ind_ + 1] = sample_weight_[prev_ind_] / 0.99;
                sum = sum + sample_weight_[i];
            }

            else {
                sum += sample_weight_[i];
            }
            
        }
        for (int i = 0; i < 50; ++i) {
            sample_weight_[i] /= sum;
            if ( i < 13) {
                sum_init += sample_weight_[i];
            }
            else {
                sum_else += sample_weight_[i];
            }
            // cout<<"sample_weight = "<<sample_weight_[i] << " + "<<i<<endl;
        }

        if (sum_init < 0.25) {
            for (int i = 0; i < 50; ++i) { 
                if ( i < 13) {
                    sample_weight_[i] = 0.25 / 13;
                }
                else {
                    sample_weight_[i] / (0.25 + sum_else);
                }
            }
        }
        prev_ind_ = replace_id + 12;

        ret = (Result)aclrtMemcpy(sample_weightBuf_, sample_weightSize_,
                            sample_weight_, sample_weightSize_, policy);
    
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy sample_weight_ to device failed.");
            return FAILED;
        }
        // cout<<"copy success!"<<endl;
    }



    if (frame_nums % 20 == 0 || now_state == HARD_NEGATIVE) {
        ret = (Result)aclrtMemcpy(origin_weight_DataBuf_, 16384,
                            origin_weight, 16384, policy);
    
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("Copy initweight to device failed.");
            return FAILED;
        }
    }

    return SUCCESS;
}

Result Optimizer::InferenceForInitialOptimizer() {
    aclmdlDataset* inferenceOutput;
    //Perform reasoning
    AtlasError ret;
    // ret = aclrtSynchronizeStream(stream_);
    ret = initial_optimizer_.Execute(stream_, 0);
    if (ret != SUCCESS) {
        ERROR_LOG("Execute model inference failed");
        return FAILED;
    }
    //Get inference output
    inferenceOutput = initial_optimizer_.GetModelOutputData();

    final_weights_.firstaddr = GetInferenceOutputItem(final_weights_.size, inferenceOutput, 0);
    // cout<<"size of final_weights = "<<final_weights_.size<<endl;

    // string modelname = "Optimizer";
    // initial_optimizer_.OutputModelResult(modelname);



    return SUCCESS;
}

Result Optimizer::InferenceForTrackingOptimizer(bool update_flag, bool get_result) {
    aclmdlDataset* inferenceOutput;

    if (update_flag == true) {
        //Perform reasoning
        AtlasError ret = tracking_optimizer_.Execute(stream_, 0);
        // if (ret != SUCCESS) {
        //     ERROR_LOG("Execute model inference failed");
        //     return FAILED;
        // }
    }

    if (get_result == true) {
        //Get inference output
        // aclrtSynchronizeStream(update_stream_);
        inferenceOutput = tracking_optimizer_.GetModelOutputData();

        new_weights_.firstaddr = GetInferenceOutputItem(new_weights_.size, inferenceOutput, 0);
        cout<<"size of final_weights = "<<final_weights_.size<<endl;

        // string modelname = "Optimizer";
        // initial_optimizer_.OutputModelResult(modelname);
    }
   
    return SUCCESS;
}

void* Optimizer::GetInferenceOutputItem(uint32_t& itemDataSize,
                                           aclmdlDataset* inferenceOutput,
                                           uint32_t idx) {

    aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(inferenceOutput, idx);
    if (dataBuffer == nullptr) {
        ERROR_LOG("Get the %dth dataset buffer from model "
        "inference output failed", idx);
        return nullptr;
    }

    void* dataBufferDev = aclGetDataBufferAddr(dataBuffer);
    if (dataBufferDev == nullptr) {
        ERROR_LOG("Get the %dth dataset buffer address "
        "from model inference output failed", idx);
        return nullptr;
    }

    size_t bufferSize = aclGetDataBufferSize(dataBuffer);
    if (bufferSize == 0) {
        ERROR_LOG("The %dth dataset buffer size of "
        "model inference output is 0", idx);
        return nullptr;
    }

    void* data = nullptr;
    if (runMode_ == ACL_HOST) {
        data = Utils::CopyDataDeviceToLocal(dataBufferDev, bufferSize);
        if (data == nullptr) {
            ERROR_LOG("Copy inference output to host failed");
            return nullptr;
        }
    }
    else {
        data = dataBufferDev;
    }

    itemDataSize = bufferSize;
    return data;
}

void Optimizer::DestroyResource()
{

    initial_optimizer_.Unload();
    initial_optimizer_.DestroyDesc();
    initial_optimizer_.DestroyInput();
    initial_optimizer_.DestroyOutput();
}
